package org.codeart.juc.pool;

import java.util.concurrent.Executors;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;
import java.util.concurrent.RecursiveTask;

/**
 * 工作窃取线程池
 */
public class WorkStealingThreadPoolDemo {

    public static int[] nums = new int[1_000_000_000];
    
    static {
        for (int i = 0; i < nums.length; i++) {
            nums[i] = (int) (Math.random() * 1000);
        }
    }

    public static class SumRecursiveTask extends RecursiveTask<Integer> {
        
        // 指定一个线程处理索引位置
        private final int start, end;
        
        // 步长
        private static final int MAX_STRIDE = 200_000_000;

        public SumRecursiveTask(int start, int end) {
            this.start = start;
            this.end = end;
        }

        @Override
        protected Integer compute() {
            int ans = 0;
            int stride = end - start;
            
            // 可以处理任务
            if (stride <= MAX_STRIDE) {
                for (int i = start; i <= end; i++) {
                    ans += nums[i];
                }
            } else {
                int mid = (start + end) >>> 1;
                // 声明为两个任务
                SumRecursiveTask left = new SumRecursiveTask(start, mid);
                SumRecursiveTask right = new SumRecursiveTask(mid + 1, end);
                
                // 分别执行两个任务
                left.fork();
                right.fork();
                
                // 等待结果
                ans += left.join() + right.join();
            }
            return ans;
        }
    }

    public static void main(String[] args) {
        // 单线程方式
        System.out.println("单线程计算数组总和");
        long start = System.currentTimeMillis();
        int ans = 0;
        for (int num : nums) {
            ans += num;
        }
        long end = System.currentTimeMillis();
        System.out.println("单线程运算结果: " + ans + " 耗时: " + (end - start) + "ms");
        
        // 多线程方式
        // 使用 ForkJoinPool 时
        // Runnable -> RecursiveAction
        // Callable<V> -> RecursiveTask<V>
        System.out.println("多线程计算数组总和");
        ForkJoinPool pool = (ForkJoinPool) Executors.newWorkStealingPool();
        long start1 = System.currentTimeMillis();
        ForkJoinTask<Integer> task = pool.submit(new SumRecursiveTask(0, nums.length - 1));
        int ans1 = task.join();
        long end1 = System.currentTimeMillis();
        System.out.println("多线程运算结果: " + ans1 + " 耗时: " + (end1 - start1) + "ms");
    }

}
